package com.liuyu.common.framework.lang.util.ratelimiter;

/**
 * 令牌桶算法
 * 系统服务作为生产者，按照指定频率向桶（容器）中添加令牌，如 QPS 为 500，每 2ms 向桶中添加一个令牌，如果桶中令牌数量达到阈值，则不再添加。
 * 请求的执行作为消费者，每个请求都需要去桶中拿取一个令牌，取到令牌则继续执行；
 * 如果桶中无令牌可取，就触发拒绝策略，可以是超时等待，也可以是直接拒绝请求，以此达到限流目的。
 * <p>
 * 1s / 限流阈值（QPS）  = 令牌添加时间间隔；
 * 桶的容量可以大于限流的阈值（做一定的冗余），令牌数量达到桶容量时，不再添加；
 *
 * @author: liuyu
 * @date: 2023/5/22 9:08
 * @version: 1
 */

import java.time.LocalDateTime;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;

/**
 * 令牌桶算法
 * 系统服务作为生产者，按照指定频率向桶（容器）中添加令牌，如 QPS 为 500，每 2ms 向桶中添加一个令牌，如果桶中令牌数量达到阈值，则不再添加。
 * 请求的执行作为消费者，每个请求都需要去桶中拿取一个令牌，取到令牌则继续执行；
 * 如果桶中无令牌可取，就触发拒绝策略，可以是超时等待，也可以是直接拒绝请求，以此达到限流目的。
 * <p>
 * 1s / 限流阈值（QPS）  = 令牌添加时间间隔；
 * 桶的容量可以大于限流的阈值（做一定的冗余），令牌数量达到桶容量时，不再添加；
 *
 * @author Saint
 */
public class TokenBucketRateLimiter<T> {

    private int qps = 2;

    /**
     * 漏桶
     */
    private volatile LinkedBlockingQueue<Character> tokenBucket = null;

    private volatile LinkedBlockingQueue<T> taskQueue = null;


    /**
     * 表示桶是否被初始化
     */
    private volatile boolean initialized = false;

    /**
     * 关联一个漏桶（一般情况为自身）
     */
    private static volatile TokenBucketRateLimiter tokenBucketRateLimiter;

    public static TokenBucketRateLimiter getTokenBucket() {
        return tokenBucketRateLimiter;
    }

    /**
     * 创建一个漏桶
     *
     * @param capacity 漏桶容量
     * @param qps      漏桶支持的QPS
     * @return
     */
    public static TokenBucketRateLimiter create(int capacity, int qps) {
        tokenBucketRateLimiter = new TokenBucketRateLimiter(capacity, qps);
        return tokenBucketRateLimiter;
    }

    private TokenBucketRateLimiter(int capacity, int qps) {
        // 漏桶只能被初始化一次
        if (!initialized) {
            initialized = true;
            this.qps = qps;
            tokenBucket = new LinkedBlockingQueue<>(capacity);
            // 初始化生产者
            initProducer();
        }
    }

    private void startTask() {
        T task = taskQueue.poll();

    }


    /**
     * 令牌桶中的令牌以固定速率加入
     */
    private void initProducer() {
        // 令牌桶初始容量为qps
        for (int i = 0; i < qps; i++) {
            tokenBucket.offer('S');
        }
        new Thread(() -> {
            while (true) {
                try {
                    TimeUnit.MILLISECONDS.sleep(1000 / qps);
                } catch (InterruptedException e) {
                    // log for exception
                    System.out.println("Exception occur! " + e);
                }
                // 以固定速率生产令牌
                tokenBucket.offer('S');
            }
        }).start();
    }

    /**
     * 获取令牌，能获取到，允许请求通过，返回true；否则返回false
     *
     * @return
     */
    public boolean tryAcquire() {
        return tokenBucket.poll() == null ? false : true;
    }

    /**
     * 令牌桶容量为2，每秒加入2个令牌；
     * 最初漏桶中有两个令牌，所以第一秒立刻获取到两个令牌；又由于令牌桶每秒加入2个令牌，所以在程序开始跑时，第一秒必回加入一个令牌（极限情况2个）
     * 所以第一秒会获取到3个令牌，第一秒之后每秒稳定获取到两个令牌，即QPS为2；
     *
     * @param args
     */
    public static void main(String[] args) {
        TokenBucketRateLimiter tokenBucket = TokenBucketRateLimiter.create(10, 2);
        int i1 = 20;
        for (int i = 0; i < i1; i++) {
            LocalDateTime now = LocalDateTime.now();
            if (tokenBucket.tryAcquire()) {
                System.out.println(i + "," + now + " pass the rate limiting");
            } else {
                System.out.println(i + "," + now + " was limited");
            }

            try {
                TimeUnit.MILLISECONDS.sleep(200);
            } catch (InterruptedException e) {
                // 日志记录
            }

        }

        System.out.println("------------");
        // 再次获取漏桶
        TokenBucketRateLimiter tokenBucket2 = TokenBucketRateLimiter.getTokenBucket();

        // 验证令牌桶只会有一个
        System.out.println("tokenBucket only one ? " + (tokenBucket == tokenBucket2));

        int i2 = 10;
        for (int i = 0; i < i2; i++) {
            LocalDateTime now = LocalDateTime.now();
            if (tokenBucket2.tryAcquire()) {
                System.out.println(i + "," + now + " pass the rate limiting");
            } else {
                System.out.println(i + "," + now + " was limited");
            }

            try {
                TimeUnit.MILLISECONDS.sleep(200);
            } catch (InterruptedException e) {
                // 日志记录
            }

        }
    }

}

